Algoritam Diffie-Hellman nad grupom tačkama eliptičke krive

import random

# Pomoćna funkcija, prošireni Euklidov algoritam
def ext_gcd(a, b):
    if b == 0:
        return (a, 1, 0)
    g, x, y = ext_gcd(b, a % b)
    return (g, y, x - a // b * y)

def gcd(a, b):
    if b == 0:
        return a
    return gcd(b, a % b)

def mod_inv(a, m):
    g, x, y = ext_gcd(a, m)
    if g != 1:
#         print("Vrednosti a i m nisu uzajamno proste!")
        return None
    else:
        return x % m

class EC_Point:
    def __init__(self, x, y, q, a):
        self.x = x
        self.y = y
        self.q = q
        self.a = a
        
    def __str__(self):
        return f'({self.x},{self.y})'
        
    def add(self, Q):
        if self.x == 0 and self.y == 0:
            return Q
        
        if Q.x == 0 and Q.y == 0:
            return self
        
        if self.x == Q.x and self.y == Q.y:
            return self.double()
        
        try:
            k = (Q.y - self.y) * mod_inv(Q.x - self.x, self.q)
            n = (-k * self.x + self.y) % self.q
        except:
            return EC_Point(0, 0, self.q, self.a)
        
        xr_ = (pow(k, 2, self.q) - self.x - Q.x) % self.q
        yr_ = (k * xr_ + n) % self.q
        
        return EC_Point(xr_, -yr_ % self.q, self.q, self.a)
        
    def double(self):
        if self.x == 0 and self.y == 0:
            return EC_Point(0, 0, self.q, self.a)
        
        try:
            k = ((3 * pow(self.x, 2, self.q) + self.a) * mod_inv(2 * self.y, self.q)) % self.q
            n = (-k * self.x + self.y) % self.q

            xr_ = (pow(k, 2, self.q) - 2 * self.x) % self.q
            yr_ = (k * xr_ + n) % self.q
        except:
            return EC_Point(0, 0, self.q, self.a)
        
        return EC_Point(xr_, -yr_ % self.q, self.q, self.a)
    
    def neg(self):
        return EC_Point(self.x, -self.y % self.q, self.q, self.a)
    
    def mul_scalar(self, n):
        if n == 0:
            return EC_Point(0, 0, self.q, self.a)
        
        if n == 1:
            return self
        
        if n % 2 == 1:
            return self.mul_scalar(n - 1).add(self)
            
        return self.mul_scalar(n // 2).double()

# y^2 = x^3 + Ax + B
class EC:
    def __init__(self, a, b, q, G):
        self.a = a
        self.b = b
        self.q = q
        self.G = G
        
    def point(self, x, y):
        return EC_Point(x, y, self.q, self.a)

class EC_Diffie_Hellman(EC):
    def __init__(self, a, b, q, gen, priv=None):
        (xg, yg) = gen
        
        G = EC_Point(xg, yg, q, a)
        
        super().__init__(a, b, q, G)
            
        if priv == None:
            priv = random.randrange(2, q)
            
        self.priv = priv
        self.pub = G.mul_scalar(priv)
        self.key = None
        
    def calculate_key(self, B_pub):
        self.key = B_pub.mul_scalar(self.priv)
        return self.key

# Primer

q = 211

# Vrednosti postavljene radi primer, promeniti na None 
priv_a = 121
priv_b = 223

a = 0
b = -4
gen = (2,2)

A = EC_Diffie_Hellman(a, b, q, gen, priv_a)
B = EC_Diffie_Hellman(a, b, q, gen, priv_b)

print(A.calculate_key(B.pub))
print(B.calculate_key(A.pub))

(111,66)
(111,66)